/**
 * This program is free software, you can redistribute it and/or modify it.
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file NMSWithMasker.h
 * \brief
*/

#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 1; // tensor num for each queue
using namespace AscendC;


template <typename T>
class KernelNMSWithMask {
public:
    __aicore__ inline KernelNMSWithMask(){};

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, const TilingData &tiling_data)
    {
        
        this->rawLength = tiling_data.totalLength;
        this->totalLength = (tiling_data.totalLength + 7) / 8 * 8;
        this->scoresLength = ((tiling_data.totalLength / 4) + 7)/ 8 * 8;
        this->outLength = ((tiling_data.totalLength / 4) + 31)/ 32 * 32;

        this->iou_threshold = tiling_data.iou_threshold;
        this->scores_threshold = tiling_data.scores_threshold;

        inputGMX.SetGlobalBuffer((__gm__ T*)x , totalLength);
        inputGMY.SetGlobalBuffer((__gm__ T*)y , scoresLength);
        outputGMZ.SetGlobalBuffer((__gm__ uint8_t*)z , outLength);//一个box对应四个输入

        pipe.InitBuffer(inputQueueX, BUFFER_NUM, totalLength * sizeof(T));
        pipe.InitBuffer(inputQueueY, BUFFER_NUM, scoresLength * sizeof(T));
        pipe.InitBuffer(outputQueueZ, BUFFER_NUM, outLength * sizeof(uint8_t));

        pipe.InitBuffer(tmpBuffer, outLength * sizeof(uint8_t));
    }


    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<T> xLocal = inputQueueX.AllocTensor<T>();
        AscendC::LocalTensor<T> yLocal = inputQueueY.AllocTensor<T>();
        AscendC::DataCopy(xLocal, inputGMX[0], totalLength);
        AscendC::DataCopy(yLocal, inputGMY[0], scoresLength);
        inputQueueX.EnQue(xLocal);
        inputQueueY.EnQue(yLocal);
    }

    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<uint8_t> zLocal = outputQueueZ.DeQue<uint8_t>();
        AscendC::DataCopy(outputGMZ[0], zLocal, static_cast<uint32_t>(outLength));
        outputQueueZ.FreeTensor(zLocal);
    }

    __aicore__ inline void IOU(LocalTensor<T>& xLocal, LocalTensor<T>& yLocal, 
            LocalTensor<uint8_t>& zLocal, LocalTensor<uint8_t>& tmpLocal){
         // -------------- IOU --------------
        // 1. 将所有的边界框按照置信度分数进行降序排序
        // 2. 选择分数最高的边界框，将其加入到保留列表中
        // 3. 计算该边界框与所有剩余边界框的IOU值
        // 4. 删除IOU值高于设定阈值的所有边界框-》label = 0
        // 5. 重复步骤2-4直到所有边界框都被处理
        for(uint32_t i = 0 ; i < (this->rawLength / 4) ; i++){
            tmpLocal.SetValue(i, 1); // Initially marked as retain
            zLocal.SetValue(i, 0); // Initialize output mask to 0
        }
         for(uint32_t i = 0 ; i < (this->rawLength / 4) ; i++){
             if(tmpLocal.GetValue(i) == 0){
                continue; //该框已被标记为不保留，跳过
            }
            if(yLocal.GetValue(i) < scores_threshold){
                tmpLocal.SetValue(i, 0);
                continue;//低于置信度阈值，跳过
            }
            zLocal.SetValue(i, 1); //标记为保留
            T x1_a = xLocal.GetValue(i * 4 + 0);
            T y1_a = xLocal.GetValue(i * 4 + 1);
            T x2_a = xLocal.GetValue(i * 4 + 2);
            T y2_a = xLocal.GetValue(i * 4 + 3);
            T area1 = (x2_a - x1_a) * (y2_a - y1_a);
            for(uint32_t j = i + 1 ; j < (rawLength / 4) ; j++){
                if(tmpLocal.GetValue(j) == 0){
                    continue; //该框已被标记为不保留，跳过
                }
                if(yLocal.GetValue(j) < scores_threshold){
                    tmpLocal.SetValue(j, 0);
                    continue;//低于置信度阈值，跳过
                }
                T x1_b = xLocal.GetValue(j * 4 + 0);
                T y1_b = xLocal.GetValue(j * 4 + 1);
                T x2_b = xLocal.GetValue(j * 4 + 2);
                T y2_b = xLocal.GetValue(j * 4 + 3);
                T area2 = (x2_b - x1_b) * (y2_b - y1_b);

                //计算交集坐标
                T x_inter_l = x1_a > x1_b ? x1_a : x1_b;
                T y_inter_l = y1_a > y1_b ? y1_a : y1_b;
                T x_inter_r = x2_a < x2_b ? x2_a : x2_b;
                T y_inter_r = y2_a < y2_b ? y2_a : y2_b;
                
                //计算交集面积
                T w = 0 > (x_inter_r  - x_inter_l) ? 0 :(x_inter_r  - x_inter_l) ;
                T h = 0 > (y_inter_r  - y_inter_l) ? 0 :(y_inter_r  - y_inter_l) ;
                T inter = w * h;

                //计算IOU
                T ovr = inter / (area1 + area2 - inter);
                if(ovr > this->iou_threshold){
                    tmpLocal.SetValue(j, 0); // Mark as discarded
                }
            }
        }
    }


    __aicore__ inline void Compute()
    {
        //1.数据预处理
        //2.进行NMS操作：
        //========================Preprocess========================
        AscendC::LocalTensor<T> xLocal = inputQueueX.DeQue<T>();
        AscendC::LocalTensor<T> yLocal = inputQueueY.DeQue<T>();
        AscendC::LocalTensor<uint8_t> zLocal = outputQueueZ.AllocTensor<uint8_t>();
        AscendC::LocalTensor<uint8_t> tmpLocal = tmpBuffer.Get<uint8_t>();//表示是否保留该框的标记，1表示保留，0表示不保留

        //========================Process========================
        IOU(xLocal, yLocal, zLocal, tmpLocal);

        //========================Postprocess========================
        outputQueueZ.EnQue<uint8_t>(zLocal);
        inputQueueX.FreeTensor(xLocal);
        inputQueueY.FreeTensor(yLocal);
    }

    __aicore__ inline void Process(){
        CopyIn();
        Compute();
        CopyOut();
    }




private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inputQueueX;
    TQue<QuePosition::VECIN, BUFFER_NUM> inputQueueY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outputQueueZ;
    TBuf<QuePosition::VECCALC> tmpBuffer;
    GlobalTensor<T> inputGMX;
    GlobalTensor<T> inputGMY;
    GlobalTensor<uint8_t> outputGMZ;

    int64_t rawLength;

    int64_t totalLength;
    int64_t scoresLength;
    int64_t outLength;
    
    float iou_threshold;
    float scores_threshold;
};
            


                
// =========================================================
//                 Kernel Entry
// =========================================================
extern "C" __global__ __aicore__ void nms_with_mask_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    KernelNMSWithMask<DTYPE_X> op; 
    op.Init(x, y, z, workspace, tiling_data); 
    op.Process();
}
